import numpy as np

from loss import cross_entropy_error


def softmax(a):
    c = np.max(a)
    exp_a= np.exp(a-c)
    sum_exp_a = np.sum(exp_a)
    y = exp_a /sum_exp_a
    return y

class Affine:
    '''用于计算仿射变换和梯度'''
    def __init__(self, W ,b):
        '''
        初始化参数 
        W:ndarray 权重矩阵, (input_size,output_size)
        b:ndarray 偏置参数,(output_size)
        '''
        self.params = [W,b]
        self.grads = [np.zeros_like(W),np.zeros_like(b)]
        self.x = None
        
    def forward(self,x):
        '''x 是 batch_size * 维度'''
        h = np.dot(x,self.params[0]) + self.params[1]
        self.x = x
        return h

    def backward(self,dout):
        dx = np.dot(dout,self.params[0].T)
        self.grads[0] = np.dot(self.x.T,dout)
        self.grads[1] = np.sum(dout,axis=0)
        return dx

class Sigmoid:
    def __init__(self):
        self.out = None

    def forward(self,x):
        out = 1/(1 + np.exp(-x))
        self.out = out

        return out
    
    def backward(self,dout):
        dx = dout * (1.0 - self.out) * self.out
        return dx
    


class ReLu:
    def __init__(self):
        self.mask = None

    def forward(self,x):
        self.mask = (x <= 0)
        out = x.copy()
        out[self.mask] = 0

        return out
        
    def backward(self,dout):
        dout[self.mask] = 0
        dx = dout
        return dx

class SoftmaxWithLoss:
    def __init__(self):
        self.loss = None
        self.y = None
        self.t = None

    def forward(self,x):
        self.t = t
        self.y = softmax(x)
        self.loss = cross_entropy_error(self.y,self.t)
        return self.loss
    
    def backward(self,dout=1):
        batch_size = self.t.shape[0]
        dx = (self.y - self.t) / batch_size

        return dx


if __name__ == "__main__":
    W = np.random.randn(2,3)
    b = np.random.randn(3)
    print(W)
    print(b)
    layer = Affine(W,b)
    x = np.random.randn(5,2)
    print(layer.forward(x))
    print(layer.backward(np.ones_like(layer.forward(x))))

                         